import torch
from torch.optim.sgd import SGD
from torch.nn import MarginRankingLoss
from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm

from model.transe import TransE
from .metrics import mrr_score, hit_score
from .saver import CkptSaver


class Trainer(object):
    def __init__(self, model:TransE, train_dataset, validate_dataset, margin=1.0, \
                 lr=1e-5, batch_size=32, weights=None, max_ckpt=2, outputs="outputs/"):
        self.optimizer = SGD(model.parameters(), lr=lr)
        self.loss_fn = MarginRankingLoss(margin=margin)
        cuda_condition = torch.cuda.is_available()
        self.device = torch.device("cuda:0" if cuda_condition else "cpu")
        self.model = model.to(self.device)
        self.saver = CkptSaver(model, max_ckpts=max_ckpt, outputs=outputs)

        if weights:
            self.model.entities_embedding.weight.data.copy_(weights["entities_embedding"])
            self.model.relations_embedding.weight.data.copy_(weights["relations_embedding"])
        
        self.train_dloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)
        self.val_dloader = DataLoader(validate_dataset, batch_size=int(batch_size**0.5), shuffle=False)

    def train(self, epoch):
        data_ietr = tqdm(self.train_dloader, desc="Epoch {}-".format(epoch), total=len(self.train_dloader))
        
        total_loss = 0.0
        for data in data_ietr:
            data = [d.to(self.device) for d in data if isinstance(d, torch.Tensor)] # 全部移到计算设备上
            positive, negtive = data[:3], data[3:]  # 拆分正负例
            p_score = self.model(positive)
            n_score = self.model(negtive)
            loss = self.loss_fn(p_score, n_score, torch.Tensor([1]).to(self.device))

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            total_loss += loss.item()
        self.saver.update(total_loss / len(self.train_dloader))

        return total_loss / len(self.train_dloader)

    def validate(self, epoch):
        data_ietr = tqdm(self.val_dloader, desc="Epoch {}-".format(epoch), total=len(self.val_dloader))
        total_mrr = 0.0
        total_hit5 = 0.0
        for data in data_ietr:
            data = [d.to(self.device) for d in data if isinstance(d, torch.Tensor)] # 全部移到计算设备上
            inputs, true_y = data[:2], data[2]
            pred_y = self.model.predict(inputs)

            total_mrr += mrr_score(pred_y, true_y).item()
            total_hit5 += hit_score(pred_y, true_y, 5).item()

        return {"MRR": total_mrr / len(self.val_dloader), "hit@5": total_hit5 / len(self.val_dloader)}
